import yaml
import av
import threading
import queue
import time
import logging
import os
from PIL import Image
import numpy as np
import cv2
from flask import Flask, make_response

# -------------------------- 1. 基础配置 --------------------------
logging.basicConfig(level=logging.INFO, format='%(asctime)s - Media - %(levelname)s - %(message)s')
logger = logging.getLogger("MediaPlayer")

IMAGE_CONFIG = {
    "no_signal": {"ref": "asset/no_signal_reference.png", "replace": "asset/placeholder.jpg"},
    "loading": {"ref": "asset/loading.jpg", "replace": "asset/loadinghnu.jpg"},
    "manuscript": {"ref": "asset/manuscript.jpg", "replace": "asset/manuscripthnu.jpg"},
    "ekko": {"ref": "asset/load_ekko.png", "replace": "asset/load_ekko.jpg"},
    "shutdown": {"ref": "asset/shutdown.png", "replace": "asset/shutdown.jpg"},
    "placeholder": "asset/placeholder.jpg"
}

HTTP_CONFIG = {
    "host": "127.0.0.1",
    "port": 9010,
    "jpeg_quality": 120,
    "route": "/media"
}


# -------------------------- 2. 预加载参考图 --------------------------
def load_reference_image(path, default_size=(800, 480)):
    try:
        img = Image.open(path).convert("RGB")
        logger.info(f"成功加载参考图：{path}")
        return img
    except Exception as e:
        logger.error(f"加载参考图失败：{path}，原因：{str(e)}")
        return Image.new("RGB", default_size, color="black")


NO_SIGNAL_REF = load_reference_image(IMAGE_CONFIG["no_signal"]["ref"])
LOAD_REF = load_reference_image(IMAGE_CONFIG["loading"]["ref"])
MANUSCRIPT_REF = load_reference_image(IMAGE_CONFIG["manuscript"]["ref"])
EKKO_REF = load_reference_image(IMAGE_CONFIG["ekko"]["ref"])
SHUTDOWN_REF = load_reference_image(IMAGE_CONFIG["shutdown"]["ref"])

LOAD_REPLACE = load_reference_image(IMAGE_CONFIG["loading"]["replace"])
MANUSCRIPT_REPLACE = load_reference_image(IMAGE_CONFIG["manuscript"]["replace"])
EKKO_REPLACE = load_reference_image(IMAGE_CONFIG["ekko"]["replace"])
SHUTDOWN_REPLACE = load_reference_image(IMAGE_CONFIG["shutdown"]["replace"])
PLACEHOLDER = load_reference_image(IMAGE_CONFIG["placeholder"])


# -------------------------- 3. Media类（保持不变） --------------------------
class MediaServer:
    def __init__(self, config, error_callback=None, status_callback=None):
        self.config = config
        self.player = config['player']
        self.url = self.player['url']

        self.FRAME_QUEUE_SIZE = 3
        self.frame_queue = queue.Queue(maxsize=self.FRAME_QUEUE_SIZE)

        self.reference_map = {
            "loading": (LOAD_REF, LOAD_REPLACE),
            "manuscript": (MANUSCRIPT_REF, MANUSCRIPT_REPLACE),
            "no_signal": (NO_SIGNAL_REF, PLACEHOLDER),
            "ekko": (EKKO_REF, EKKO_REPLACE),
            "shutdown": (SHUTDOWN_REF, SHUTDOWN_REPLACE)
        }

        self.playing = False
        self.thread = None  # 帧读取线程
        self.update_thread = None  # 持续帧更新线程
        self.container = None
        self.stream = None

        self.last_frame_received = time.time()
        self.frame_count = 0
        self.latest_frame = None  # 实时缓存最新帧
        self.update_interval = 0.05  # 帧更新间隔（20帧/秒）

        self.error_callback = error_callback
        self.status_callback = status_callback

        self.reconnecting = False
        self.reconnect_interval = 5
        self.reconnect_attempts = 0

    def start(self):
        if self.playing:
            logger.warning("视频流已在运行，无需重复启动")
            return

        self.playing = True
        self.reconnecting = False
        self.reconnect_attempts = 0
        self._clear_queue()

        if self.status_callback:
            self.status_callback("connecting")

        # 启动帧读取线程
        self.thread = threading.Thread(target=self._read_frames, daemon=True)
        self.thread.start()
        # 启动后台持续帧更新线程
        self.update_thread = threading.Thread(target=self._update_latest_frame_loop, daemon=True)
        self.update_thread.start()

        logger.info(f"视频流线程+帧更新线程已启动，目标地址：{self.url}")

    def stop(self):
        self.playing = False
        self.reconnecting = False
        self._cleanup_resources()

        # 停止帧更新线程
        if self.update_thread and self.update_thread.is_alive():
            self.update_thread.join(timeout=1)
        # 停止帧读取线程
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=2)

        if self.status_callback:
            self.status_callback("stopped")
        logger.info("视频流已停止，所有资源已清理")

    def get_latest_frame(self):
        """获取最新PIL帧（被动调用，优先用缓存）"""
        if self.latest_frame is not None:
            return self.latest_frame.copy()

        # 兜底：缓存未初始化时，从队列取帧
        frame = None
        try:
            while not self.frame_queue.empty():
                frame = self.frame_queue.get_nowait()
            self.last_frame_received = time.time()
        except queue.Empty:
            if time.time() - self.last_frame_received > 5:
                logger.warning("5秒未收到新帧，返回无信号图")
                frame = PLACEHOLDER.copy()
            else:
                frame = PLACEHOLDER.copy()

        # 帧匹配与替换
        if frame is not None:
            for ref_name, (ref_img, replace_img) in self.reference_map.items():
                if self.is_frame_matched(frame, ref_img):
                    logger.debug(f"匹配到{ref_name}帧，使用替换图")
                    frame = replace_img.copy()
                    break

        return frame

    def get_latest_jpeg(self):
        """获取最新JPEG字节流（供HTTP调用）"""
        frame = self.latest_frame or PLACEHOLDER.copy()
        ret, jpeg_data = cv2.imencode(
            '.jpg',
            cv2.cvtColor(np.array(frame), cv2.COLOR_RGB2BGR),
            [int(cv2.IMWRITE_JPEG_QUALITY), HTTP_CONFIG["jpeg_quality"]]
        )
        return jpeg_data.tobytes() if ret else b''

    def _update_latest_frame_loop(self):
        logger.info("后台帧更新线程已启动，每0.05秒刷新一次最新帧")
        while self.playing:
            try:
                # 循环从队列取最新帧（清空旧帧，只留最新）
                frame = None
                while not self.frame_queue.empty():
                    frame = self.frame_queue.get_nowait()
                    self.last_frame_received = time.time()
                    self.frame_count += 1

                # 帧匹配与替换
                if frame is None:
                    # 无新帧时，用缓存帧或无信号图兜底
                    if time.time() - self.last_frame_received > 5:
                        frame = PLACEHOLDER.copy()
                    else:
                        frame = self.latest_frame.copy() if self.latest_frame else PLACEHOLDER.copy()
                else:
                    # 有新帧时，执行替换逻辑
                    for ref_name, (ref_img, replace_img) in self.reference_map.items():
                        if self.is_frame_matched(frame, ref_img):
                            logger.debug(f"后台更新：匹配到{ref_name}帧，使用替换图")
                            frame = replace_img.copy()
                            break

                # 更新缓存帧
                self.latest_frame = frame.copy()

                # 控制更新频率
                time.sleep(self.update_interval)
            except Exception as e:
                logger.error(f"后台帧更新线程出错：{str(e)}，1秒后重试")
                time.sleep(1)

    def is_frame_matched(self, frame, ref_img, threshold=0.8):
        if frame is None or ref_img is None:
            return False

        def pil2cv(pil_img):
            return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

        try:
            frame_cv = pil2cv(frame)
            ref_cv = pil2cv(ref_img)
            ref_cv_resized = cv2.resize(ref_cv, (frame_cv.shape[1], frame_cv.shape[0]))

            frame_gray = cv2.cvtColor(frame_cv, cv2.COLOR_BGR2GRAY)
            ref_gray = cv2.cvtColor(ref_cv_resized, cv2.COLOR_BGR2GRAY)
            result = cv2.matchTemplate(frame_gray, ref_gray, cv2.TM_CCOEFF_NORMED)
            max_val = cv2.minMaxLoc(result)[1]

            return max_val >= threshold
        except Exception as e:
            logger.error(f"帧匹配失败：{str(e)}")
            return False

    def _read_frames(self):
        while self.playing:
            try:
                self._open_stream()
                self._frame_read_loop()
            except Exception as e:
                self._cleanup_resources()
                if self.playing and not self.reconnecting:
                    self._start_reconnecting()
                    time.sleep(self.reconnect_interval)

    def _open_stream(self):
        if self.status_callback:
            self.status_callback("connecting")
        logger.info(f"尝试连接视频流：{self.url}")

        self.container = av.open(
            self.url,
            mode='r',
            options={
                'rtsp_transport': 'tcp',
                'stimeout': '5000000',
                'buffer_size': '262144',
                'flags': 'low_delay',
                'framedrop': '1'
            }
        )

        self.stream = next((s for s in self.container.streams if s.type == 'video'), None)
        if not self.stream:
            raise ValueError("未找到视频流轨道")

        self.reconnecting = False
        self.reconnect_attempts = 0
        if self.status_callback:
            self.status_callback("connected")
        logger.info("视频流连接成功，开始读取帧")

    def _frame_read_loop(self):
        for packet in self.container.demux(self.stream):
            if not self.playing or self.reconnecting:
                break
            for frame in packet.decode():
                if not self.playing:
                    break
                pil_frame = frame.to_image().convert("RGB")
                # 只往队列塞帧，不做额外处理
                if self.frame_queue.full():
                    self.frame_queue.get_nowait()
                self.frame_queue.put(pil_frame, block=False)

    def _start_reconnecting(self):
        self.reconnecting = True
        self.reconnect_attempts += 1
        if self.status_callback:
            self.status_callback(f"reconnecting({self.reconnect_attempts})")
        logger.warning(f"视频流断开，第{self.reconnect_attempts}次重连（间隔{self.reconnect_interval}秒）")

    def _clear_queue(self):
        while not self.frame_queue.empty():
            self.frame_queue.get_nowait()

    def _cleanup_resources(self):
        if self.container:
            try:
                self.container.close()
            except Exception as e:
                logger.error(f"关闭流失败：{str(e)}")
        self.container = None
        self.stream = None


# -------------------------- 4. HTTP服务（只保留取帧接口） --------------------------
def create_http_server(media):
    app = Flask(__name__)

    # 只保留获取视频帧的接口（供C++或其他服务拉取帧数据）
    @app.route(HTTP_CONFIG["route"], methods=["GET"])
    def get_media_frame():
        jpeg_data = media.get_latest_jpeg()
        response = make_response(jpeg_data)
        response.headers["Content-Type"] = "image/jpeg"
        response.headers["Content-Length"] = str(len(jpeg_data))
        return response

    # 启动HTTP服务线程（后台运行，不影响C++控制）
    threading.Thread(
        target=app.run,
        args=(HTTP_CONFIG["host"], HTTP_CONFIG["port"]),
        kwargs={"threaded": True, "use_reloader": False},
        daemon=True
    ).start()

    logger.info(f"HTTP服务已启动，仅保留取帧接口：")
    logger.info(f"获取视频帧：http://{HTTP_CONFIG['host']}:{HTTP_CONFIG['port']}{HTTP_CONFIG['route']}")
    return app


# -------------------------- 5. 实例化media_server（供C++调用） --------------------------
# 加载配置文件
try:
    with open("./configs/player.yaml", "r", encoding="utf-8") as f:
        global_config = yaml.safe_load(f)
    print("MediaServer：成功加载配置文件 configs/player.yaml")
except Exception as e:
    print(f"MediaServer：加载配置文件失败，程序异常：{str(e)}")
    raise


# 状态回调函数（可选，用于C++捕获状态日志）
def media_status_callback(status):
    if status.startswith("reconnecting("):
        n = status.split("(")[1].split(")")[0]
        print(f"MediaServer：🔄 第{n}次尝试重连视频流...")
    elif status == "connecting":
        print("MediaServer：🔄 正在连接视频流...")
    elif status == "connected":
        print("MediaServer：✅ 视频流连接成功")
    elif status == "stopped":
        print("MediaServer：⏹️ 视频流已停止")
    else:
        print(f"MediaServer：📌 状态更新：{status}")


# 实例化MediaServer（全局变量，供C++通过Python API调用）
media_server = MediaServer(
    config=global_config,
    status_callback=media_status_callback
)

# 启动HTTP取帧服务（仅用于提供帧数据，不处理启动/停止）
http_app = create_http_server(media_server)
print("MediaServer：初始化完成，等待C++调用start()启动服务")


# -------------------------- 6. 供C++调用的接口（简化主函数） --------------------------
if __name__ == "__main__":
    try:
        print("MediaServer：当前为独立运行模式（仅用于测试）")
        print("提示：在C++中，通过调用media_server.start()启动，media_server.stop()停止")
        # 如需测试，可手动调用start()
        # media_server.start()
        while True:
            time.sleep(0.01)
    except KeyboardInterrupt:
        print("MediaServer：检测到退出指令")
    finally:
        media_server.stop()
        print("MediaServer：程序已退出")